from torch.utils.data import Dataset, DataLoader
import torch
import mxnet as mx
import pickle
import numpy as np
from tqdm import tqdm
import os


class VerificationDataset(Dataset):
    """
       LFW: 12K
       AgeDB-30: 12K
       CFP-FP: 14K
    """

    def __init__(self, path):
        self.path = [os.path.join(path, 'lfw.bin'), os.path.join(path, 'agedb_30.bin'),
                     os.path.join(path, 'cfp_fp.bin')]
        self.data = []
        self.label = []
        for cur in self.path:
            with open(cur, 'rb') as f:
                data, label = pickle.load(f, encoding='bytes')
                self.data.extend(data)
                self.label.extend(label)
        self.imgs = len(self.data)
        self.people = len(self.label)

    def __len__(self):
        return self.imgs

    def __getitem__(self, idx):
        img_ = mx.image.imdecode(self.data[idx]).asnumpy()
        img_ = np.transpose(img_, axes=(2, 0, 1))
        return torch.tensor(img_, dtype=torch.float32), torch.tensor([self.label[idx // 2]])


if __name__ == '__main__':
    per_batch_size = 10

    train_dataset = VerificationDataset('/home/data1/worm/datas/faces_emore')
    train_loader = DataLoader(dataset=train_dataset, batch_size=per_batch_size)
    for output in tqdm(train_loader):
        pass
